import sys, os
import json
import csv
import numpy as np
from PIL import Image
from tqdm import tqdm

def crop_resize_save(image_path, save_path, resize=(224, 224), crop_threshold=25):
    """
    打开 image_path 的图像，裁剪掉黑边，然后 resize 到目标尺寸，
    并保存为 PNG 格式到 save_path。
    
    参数：
        image_path (str): 输入图像文件路径。
        save_path (str): 输出图像保存路径（PNG 格式）。
        resize (tuple): 目标尺寸，默认 (224, 224)。
        crop_threshold (int): 裁剪时的黑边阈值，默认 25。
    
    返回：
        dict: 包含裁剪信息的字典：
              {
                  'crop_box': (top, bottom, left, right),
                  'original_size': (height, width),
                  'cropped_size': (cropped_height, cropped_width)
              }
    """
    image = Image.open(image_path)
    if image_path.lower().endswith('.ppm'):
        image = image.convert("RGB")
    image_array = np.array(image)
    original_size = [int(image_array.shape[0]), int(image_array.shape[1])]
    
    # 使用带信息的裁剪函数
    cropped_array, crop_box = crop_black_border(image_array, threshold=crop_threshold,image_name=os.path.basename(save_path))
    cropped_size = [int(cropped_array.shape[0]), int(cropped_array.shape[1])]
    
    # 转换为 Image 对象进行 resize 操作
    cropped_image = Image.fromarray(cropped_array)
    try:
        resample_filter = Image.Resampling.LANCZOS
    except AttributeError:
        resample_filter = Image.LANCZOS
    resized_image = cropped_image.resize(resize, resample_filter)
    resized_image.save(save_path, format="PNG")
    
    # 确保 crop_box 内的值也是 Python 原生类型
    crop_box = [int(x) for x in crop_box]
    
    info = {
        'crop_box': crop_box,
        'original_size': original_size,
        'cropped_size':cropped_size
    }
    return info

def crop_black_border(image_array, threshold=25,image_name=""):
    """
    裁剪图像的黑边部分，并返回裁剪后的数组和裁剪框信息 (top, bottom, left, right)。
    
    参数：
        image_array (np.array): 输入图像数组。
        threshold (int): 裁剪阈值，默认 25。
    
    返回：
        tuple: (裁剪后的图像数组，裁剪框信息)
        
    """
    if image_array.ndim == 3:
        gray = np.mean(image_array, axis=2)
    else:
        gray = image_array

    rows_black = np.all(gray <= threshold, axis=1)
    cols_black = np.all(gray <= threshold, axis=0)
    non_black_rows = np.where(~rows_black)[0]
    non_black_cols = np.where(~cols_black)[0]

    if non_black_rows.size == 0 or non_black_cols.size == 0:
        return image_array, (0, image_array.shape[0], 0, image_array.shape[1])

    top, bottom = non_black_rows[0], non_black_rows[-1] + 1
    left, right = non_black_cols[0], non_black_cols[-1] + 1

    original_height, original_width = image_array.shape[0], image_array.shape[1]
    cropped_height = bottom - top
    cropped_width = right - left

    # 检查裁剪后尺寸是否满足要求
    if cropped_width < 0.5 * original_width or cropped_height < 0.7 * original_height:
        # 如果裁剪后尺寸过小，返回原图和原始尺寸的裁剪框
        return image_array, (0, image_array.shape[0], 0, image_array.shape[1])
        raise ValueError(
            f"Cropped image {image_name} dimensions are too small: cropped_width={cropped_width} (min required {0.5 * original_width}), "
            f"cropped_height={cropped_height} (min required {0.7 * original_height})."
        )
    return image_array[top:bottom, left:right], (top, bottom, left, right)


def crop_resize_save_with_info(image_path, save_path, info, resize=(224, 224)):
    """
    打开 image_path 的图像，裁剪掉黑边后返回裁剪信息，并将 resize 后的图像保存为 PNG 格式到 save_path。
    
    参数：
        image_path (str): 输入图像文件路径。
        info: 主图像的信息，包含原图尺寸和裁剪框 (top, bottom, left, right)。
        save_path (str): 输出图像保存路径（PNG 格式）。
        resize (tuple): 目标尺寸，默认 (224, 224)。
        crop_threshold (int): 裁剪时黑边阈值，默认 25。
    
    返回：
        None
    """
    # 打开图像并检查尺寸是否匹配
    image = Image.open(image_path)
    image_array = np.array(image)

    # 从 info 中提取原图尺寸和裁剪框
    original_height, original_width = image_array.shape[0], image_array.shape[1]
    top, bottom, left, right = info['crop_box']

    # 如果原图的尺寸与 info 提供的尺寸不同，抛出错误
    if original_height != info['original_size'][0] or original_width != info['original_size'][1]:
        raise ValueError(f"Image size mismatch: {original_height}x{original_width} != "
                         f"{info['original_size'][0]}x{info['original_size'][1]}.")

    # 按照提供的裁剪框进行裁剪
    cropped_image_array = image_array[top:bottom, left:right]

    # 如果裁剪后图像大小小于目标尺寸，则抛出异常
    cropped_height, cropped_width = cropped_image_array.shape[0], cropped_image_array.shape[1]
    if cropped_width < 0.6 * original_width or cropped_height < 0.7 * original_height:
        raise ValueError(f"Cropped image dimensions are too small: cropped_width={cropped_width} "
                         f"(min required {0.6 * original_width}), cropped_height={cropped_height} "
                         f"(min required {0.7 * original_height}).")

    # 进行 resize 操作
    cropped_image = Image.fromarray(cropped_image_array)
    try:
        resample_filter = Image.Resampling.LANCZOS
    except AttributeError:
        resample_filter = Image.LANCZOS
    resized_image = cropped_image.resize(resize, resample_filter)

    # 保存为 PNG 格式
    resized_image.save(save_path, format="PNG")
def resize_save(image_path, save_path, resize):
    """
    打开 image_path 对应的图片，调整大小为 resize 指定的尺寸，
    然后保存到 save_path（PNG 格式）。

    参数：
        image_path (str): 输入图像路径。
        save_path (str): 输出图像保存路径。
        resize (tuple): 目标尺寸 (宽，高)。
    """
    image = Image.open(image_path)
    try:
        resample_filter = Image.Resampling.LANCZOS
    except AttributeError:
        resample_filter = Image.LANCZOS
    resized_image = image.resize(resize, resample_filter)
    resized_image.save(save_path, format="PNG")

def parse_csv(csv_path):
    """
    读取 CSV 文件，并返回一个字典，格式如下：
    
    {
       image_id: "<diagnosis text>"
    }
    
    其中，文本描述的生成规则为：
      - 对于 adjudicated_dr_grade（或 level，ICDR 5 级）:
           0 -> normal
           1 -> Mild Diabetic Retinopathy
           2 -> Moderate Diabetic Retinopathy
           3 -> Severe Diabetic Retinopathy
           4 -> Proliferative Diabetic Retinopathy
           
      - 如果 CSV 文件中既没有 "level" 也没有 "DR_grade" 列，则直接抛出异常。
      
    示例 CSV 文件内容：
        ,image,quality,DR_grade
        0,1_right.jpeg,1,0
    """
    dr_mapping = {
        "0": "normal",
        "1": "Mild Diabetic Retinopathy",
        "2": "Moderate Diabetic Retinopathy",
        "3": "Severe Diabetic Retinopathy",
        "4": "Proliferative Diabetic Retinopathy"
    }
    result = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames or []
        # 检查使用哪一列来表示 DR 等级
        if "level" in fieldnames:
            grade_key = "level"
        elif "DR_grade" in fieldnames:
            grade_key = "DR_grade"
        else:
            raise ValueError("CSV file must contain either 'level' or 'DR_grade' column.")
        
        for row in reader:
            image_id = row.get("image", "").strip()
            dr_val = row.get(grade_key, "").strip()
            if not image_id or not dr_val:
                continue
            result[image_id] = dr_mapping.get(dr_val, dr_val)
    return result


def merge_data_dict(train_official, test_custom):
    """
    合并两个由 CSV 解析后的数据字典：
      - 如果 test_custom 中的某个 image_name 在 train_official 中不存在，则将该键值对加入 train_official。
    返回合并后的字典。
    """
    merged = train_official.copy()
    for k, v in test_custom.items():
        if k not in merged:
            merged[k] = v
    return merged

def gather_data(data_path, tar_path, prefix='eyepacs_', resize=(224,224)):
    """
    对新数据集进行预处理：
      - 数据集根目录下要求包含：
            trainLabels/trainLabels.csv
            Label_EyeQ_test.csv
            以及存放图像的 train 和 test 文件夹。
      - 分别解析 trainLabels 和测试 CSV，生成的标签字典 key 为图像文件名（含扩展名），值为诊断描述，
        然后将二者合并。
      - 遍历 train 和 test 文件夹，对每张图像调用 crop_resize_save 进行裁剪与 resize，
        并生成新的图像名称，格式为 "{prefix}_{原始文件名}.png"（保留原始文件名，仅修改后缀并加前缀）。
      - 同时生成一个 split 字典，记录各 split 下的处理后图像名称。
      - 最终生成的标注格式为：
            {
                new_image_name: {
                    "image_name": new_image_name,
                    "image_path": "images/{new_image_name}",
                    "diagnosis": {
                        "text": <diagnosis text>
                    }
                },
                ...
            }
      - 将标注信息保存到 tar_path/annotations.json，
        将 split 信息保存到 tar_path/split/official.json。
    
    参数：
        data_path (str): 新数据集根目录，要求包含 IMAGES 文件夹和相应 CSV 文件。
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 "eyepacs_"
        resize (tuple): 目标尺寸，默认为 (224,224)
    
    返回：
        dict: 标注信息字典
    """
    # 创建目标目录
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, "images")
    os.makedirs(images_out_dir, exist_ok=True)
    split_dir = os.path.join(tar_path, "split")
    os.makedirs(split_dir, exist_ok=True)
    
    # 解析 CSV 文件生成标签字典
    train_csv_path = os.path.join(data_path, "trainLabels", "trainLabels.csv")
    test_csv_path = os.path.join(data_path, "Label_EyeQ_test.csv")
    train_official = parse_csv(train_csv_path)
    test_custom = parse_csv(test_csv_path)
    label_dict = merge_data_dict(train_official, test_custom)
    
    final_annotations = {}
    official_split = {"train": [], "test": []}
    split_list = ["train", "test"]
    
    # 遍历每个 split 文件夹
    for split in split_list:
        image_src_dir = os.path.join(data_path, split)
        image_list=sorted(os.listdir(image_src_dir))
        total_image_num=len(image_list)
        for image_file in tqdm(image_list,desc=f'{split} process',unit='images',total=total_image_num):
            if not image_file.lower().endswith(('.png','.jpg','.jpeg','.bmp','.ppm')):
                continue
            # 保留原始文件名（不含扩展名），仅加前缀并统一保存为 PNG 格式
            base_name = os.path.splitext(image_file)[0]
            new_image_name = f"{prefix}_{base_name}.png"
            dest_image_path = os.path.join(images_out_dir, new_image_name)
            # 进行图像预处理：裁剪、resize 并保存为 PNG
            crop_info=crop_resize_save(
                image_path = os.path.join(image_src_dir, image_file),
                save_path = dest_image_path,
                resize = resize,
                crop_threshold = 15
            )
            # 获取对应的诊断描述（以原始文件名作为 key）
            diag_text = label_dict.get(image_file, "")
            final_annotations[new_image_name] = {
                "image_name": new_image_name,
                "image_path": os.path.join("images", new_image_name),
                'original_path':os.path.relpath(os.path.join( os.path.join(image_src_dir, image_file)),data_path),
                'crop_info':crop_info,
                "diagnosis": {
                    "text": diag_text
                }
            }
            official_split[split].append(new_image_name)
    
    # 保存标注信息
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(final_annotations, f, indent=4)
    
    # 保存分割信息
    official_split_path = os.path.join(split_dir, "official.json")
    with open(official_split_path, "w", encoding="utf-8") as f:
        json.dump(official_split, f, indent=4)
    
    return final_annotations

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Eyepacs 数据集预处理")
    parser.add_argument("--data_path", type=str, default="../../Dataset/public_dataset/EyePACS",
                        help="新数据集根目录，要求其中包含 trainLabels, Label_EyeQ_test CSV 文件以及 train 和 test 图像文件夹")
    parser.add_argument("--tar_path", type=str, default="../../Dataset/public_processed/EyePACS",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="eyepacs_",
                        help="处理后图片名称的前缀，默认为 'eyepacs_'")
    parser.add_argument("--resize", type=int, nargs=2, default=[224,224],
                        help="目标尺寸，默认为 224 224")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
    print("Preprocessing completed. Annotations saved.")
